#include "grate_test_utils.h"
#include "grate/antidictionary.h"
#include "grate/calculator.h"

using growth_rate::LexminAntiDictionary;
using growth_rate::LexminTrie;
using growth_rate::Calculator;

namespace {

class TestAntiDictionary : public LexminAntiDictionary {
 public:
  enum ADType {
    Only_AA,
    Only_AB,
    Both_AA_and_AB
  };

  TestAntiDictionary(uint8_t alphabet_size, ADType type)
      : LexminAntiDictionary(kTrieCapacity, alphabet_size),
        type_(type) { }

 private:
  bool construct() {
    if (type_ == Only_AA || type_ == Both_AA_and_AB) {
      add_lexmin_word("aa");
      EXPECT_FALSE(has_factor_forbidden_by_trie(""));
      EXPECT_TRUE(has_factor_forbidden_by_trie("baab"));
      EXPECT_TRUE(has_factor_forbidden_by_trie("abb"));
    }
    if (type_ == Only_AB || type_ == Both_AA_and_AB) {
      add_lexmin_word("ab");
      EXPECT_FALSE(has_factor_forbidden_by_trie("a"));
      if (type_ != Both_AA_and_AB)
        EXPECT_FALSE(has_factor_forbidden_by_trie("aaaaa"));
      EXPECT_TRUE(has_factor_forbidden_by_trie("bbbabb"));
      EXPECT_TRUE(has_factor_forbidden_by_trie("abbb"));
    }
    return true;
  }

  ADType type_;
  static const size_t kTrieCapacity = 100;
};

TEST(AntiDictionaryTest, TrieContents) {
  {
    TestAntiDictionary TAD(3, TestAntiDictionary::Only_AA);
    LexminTrie *trie = TAD.build_trie();
    EXPECT_EQ("[(1, -2, -2),\n"
              "(-1, -2, -2)]", trie->dump_transitions());
    delete trie;
  }

  {
    TestAntiDictionary TAD(3, TestAntiDictionary::Only_AB);
    LexminTrie *trie = TAD.build_trie();
    EXPECT_EQ("[(1, -2, -2),\n"
              "(-2, -1, -2)]", trie->dump_transitions());
    delete trie;
  }

  {
    TestAntiDictionary TAD(3, TestAntiDictionary::Both_AA_and_AB);
    LexminTrie *trie = TAD.build_trie();
    EXPECT_EQ("[(1, -2, -2),\n"
              "(-1, -1, -2)]", trie->dump_transitions());
    delete trie;
  }
}

TEST(AntiDictionaryTest, GrowthRate) {
  {
    const double eps = 1e-4;

    {
      TestAntiDictionary *TAD =
        new TestAntiDictionary(3, TestAntiDictionary::Only_AA);
      EXPECT_NEAR(2.0, Calculator::get_language_growth_rate(TAD), eps);
    }

    {
      TestAntiDictionary *TAD =
        new TestAntiDictionary(3, TestAntiDictionary::Only_AB);
      EXPECT_NEAR(1.0, Calculator::get_language_growth_rate(TAD), eps);
    }

    {
      TestAntiDictionary *TAD =
        new TestAntiDictionary(3, TestAntiDictionary::Both_AA_and_AB);
      EXPECT_NEAR(0, Calculator::get_language_growth_rate(TAD), eps);
    }
  }
}

}  // namespace
